from core.utils.utils import get_time
from typing import List
from gensim.models import KeyedVectors
from core.utils import logging
from core.utils.utils import get_time

class GloveWord2Vector:
    '''
    It's used to transform word to glove embedding.

    Args:
        vocab_file(str): vocab file with format of word2vec
    '''
    @get_time('load_word2vec_time', True)
    def __init__(self, vocab_file: str) -> None:
        logging.info('start to load word2vec')
        # TODO: 让加载文件的速度更快
        self.glove_model = KeyedVectors.load_word2vec_format(vocab_file)
        logging.info('load word2vec finished')
    
    def _get_vector(self, token: str, lower_case: bool = True):
        '''get word embedding by token.if this token is not in glove_model keys, use `,` replace it.'''
        token = token.lower() if lower_case and isinstance(token, str) else token
        # TODO:想一想把没有的词换成逗号的做法好不好，有没有其他更好的办法
        return self.glove_model[token] if token in self.glove_model else self.glove_model[',']

    def __call__(self, tokens: List[str], lower_case: bool = True):
        embedding = [self._get_vector(token, lower_case) for token in tokens]
        return embedding